package com.example.projectmanagement.config;

import com.example.projectmanagement.exception.SensitiveXssContentException;
import com.example.projectmanagement.util.XssUtil;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import org.springframework.util.StreamUtils;

import java.io.*;
import java.nio.charset.StandardCharsets;

/**
 * XSS防护的HttpServletRequest包装器
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    
    private byte[] rawData;
    private final HttpServletRequest request;
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    public XssHttpServletRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        this.request = request;
        // 读取原始请求体
        rawData = StreamUtils.copyToByteArray(request.getInputStream());
    }
    
    @Override
    public ServletInputStream getInputStream() throws IOException {
        // 解析JSON并清理XSS
        String requestBody = new String(rawData, StandardCharsets.UTF_8);
        
        // 检测是否包含敏感的XSS内容
        if (XssUtil.containsSensitiveXssContent(requestBody)) {
            throw new SensitiveXssContentException("存在不合规字符，请修改后提交");
        }
        
        if (isJsonRequest()) {
            try {
                JsonNode jsonNode = objectMapper.readTree(requestBody);
                sanitizeJsonNode(jsonNode);
                requestBody = objectMapper.writeValueAsString(jsonNode);
            } catch (SensitiveXssContentException e) {
                // 重新抛出XSS异常
                throw e;
            } catch (Exception e) {
                // 如果不是有效的JSON，保持原始内容
            }
        } else {
            // 对于非JSON请求，清理文本内容
            requestBody = XssUtil.sanitizeText(requestBody);
        }
        
        byte[] cleanedData = requestBody.getBytes(StandardCharsets.UTF_8);
        
        return new ServletInputStream() {
            private final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(cleanedData);
            
            @Override
            public boolean isFinished() {
                return byteArrayInputStream.available() == 0;
            }
            
            @Override
            public boolean isReady() {
                return true;
            }
            
            @Override
            public void setReadListener(ReadListener readListener) {
                // 不需要实现
            }
            
            @Override
            public int read() throws IOException {
                return byteArrayInputStream.read();
            }
        };
    }
    
    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }
    
    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        if (value != null) {
            // Skip XSS validation for multipart/form-data parameters to prevent issues with form processing
            String contentType = request.getContentType();
            if (contentType != null && (contentType.contains("multipart/form-data") || contentType.contains("form-data"))) {
                return value;
            }
            // 检测是否包含敏感的XSS内容
            if (XssUtil.containsSensitiveXssContent(value)) {
                throw new SensitiveXssContentException("存在不合规字符，请修改后提交");
            }
            return XssUtil.sanitizeText(value);
        }
        return value;
    }
    
    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values != null) {
            // Skip XSS validation for multipart/form-data parameters to prevent issues with form processing
            String contentType = request.getContentType();
            if (contentType != null && (contentType.contains("multipart/form-data") || contentType.contains("form-data"))) {
                return values;
            }
            for (int i = 0; i < values.length; i++) {
                // 检测是否包含敏感的XSS内容
                if (XssUtil.containsSensitiveXssContent(values[i])) {
                    throw new SensitiveXssContentException("存在不合规字符，请修改后提交");
                }
                values[i] = XssUtil.sanitizeText(values[i]);
            }
        }
        return values;
    }
    
    /**
     * 检查是否为JSON请求
     */
    private boolean isJsonRequest() {
        String contentType = request.getContentType();
        return contentType != null && contentType.contains("application/json");
    }
    
    /**
     * 递归清理JSON节点中的XSS
     */
    private void sanitizeJsonNode(JsonNode node) {
        if (node.isObject()) {
            ObjectNode objectNode = (ObjectNode) node;
            node.fieldNames().forEachRemaining(fieldName -> {
                JsonNode childNode = objectNode.get(fieldName);
                if (childNode.isTextual()) {
                    String textValue = childNode.asText();
                    // 检测是否包含敏感的XSS内容
                    if (XssUtil.containsSensitiveXssContent(textValue)) {
                        throw new SensitiveXssContentException("存在不合规字符，请修改后提交");
                    }
                    objectNode.put(fieldName, XssUtil.sanitizeText(textValue));
                } else {
                    sanitizeJsonNode(childNode);
                }
            });
        } else if (node.isArray()) {
            for (JsonNode childNode : node) {
                sanitizeJsonNode(childNode);
            }
        }
    }
}